#include <iostream>
#include <string>
using namespace std;
#include <assert.h>

#define __STDC_CONSTANT_MACROS

extern "C" {
#include "libavcodec/avcodec.h"
#include "libavformat/avformat.h"
#include "libswscale/swscale.h"
#include "libswresample/swresample.h"
};

#include "SDL.h"
#include "SDL_thread.h"

#include "PacketQueue.h"

#define WRITE_PCM 0

#define SDL_AUDIO_BUFFER_SIZE 1024
#define MAX_AUDIO_FRAME_SIZE 192000

#define FF_REFRESH_EVENT (SDL_USEREVENT)
#define FF_QUIT_EVENT (SDL_USEREVENT + 1)

#define ROUND(X)  ((int)((X)+0.5))

#define PRINT_RETURN_1(x) do{printf("ERROR: "#x"\n"); getchar(); return 1;}while(false)

typedef struct AudioOutSpec {
    uint64_t channelLayout;
    int channels;
    AVSampleFormat sampleFmt;
    int sampleRate;
    int bufferSize;
    int samplesPerChannel;
    void (SDLCALL *pfnCallback)(void *userdata, Uint8 *stream, int len);

    AudioOutSpec() {
        channelLayout = 0;
        channels = 0;
        sampleFmt = AV_SAMPLE_FMT_NONE;
        sampleRate = 0;
        bufferSize = 0;
        samplesPerChannel = 0;
        pfnCallback = NULL;
    }
} AudioOutSpec;

typedef struct VideoPicture {
    SDL_Overlay *pBmp;
    int width;
    int height;
    int allocated;

    VideoPicture() {
        pBmp = NULL;
        width = 0;
        height = 0;
        allocated = 0;
    }
} VideoPicture;

typedef struct VideoInfo {
    AVFormatContext *pFmtCtx;
    // video
    int videoStream;
    AVStream *pVideoStream;
    AVCodecContext *pVideoCodecCtx;
    PacketQueue videoQueue;
    struct SwsContext *pSwsCtx;

    // audio
    int audioStream;
    AVStream *pAudioStream;
    AVCodecContext *pAudioCodecCtx;
    PacketQueue audioQueue;
    struct SwrContext *pSwrCtx;
    AudioOutSpec audioOutSpec;

    // picture
    PacketQueue<VideoPicture> picQueue;
    SDL_Surface *pScreen;
	SDL_mutex *pScreenMutex;

    // thread
    SDL_Thread *pSplitPacketThread;
    SDL_Thread *pVideoDecodeThread;

    char filename[1024];
    bool quit;

    VideoInfo() {
        pFmtCtx = NULL;
        videoStream = -1;
        pVideoStream = NULL;
        pVideoCodecCtx = NULL;
        pSwsCtx = NULL;
        audioStream = -1;
        pAudioStream = NULL;
        pAudioCodecCtx = NULL;
        pSwrCtx = NULL;
        pSplitPacketThread = NULL;
        pVideoDecodeThread = NULL;
        pScreen = NULL;
        quit = false;
        filename[0] = 0;
		pScreenMutex = NULL;
    }
} VideoInfo;




#if WRITE_PCM
FILE *pcmfile;
#endif

int DecodeAudioFrame(VideoInfo *pVideoInfo, uint8_t *audio_buf, int buf_size) {
    static AVPacket pkt;
    static int audio_pkt_size = 0;
    static AVFrame frame;

    int len1;

    while(!pVideoInfo->quit) {
        while(audio_pkt_size > 0) {
            int got_frame = 0;
            len1 = avcodec_decode_audio4(pVideoInfo->pAudioCodecCtx, &frame, &got_frame, &pkt);

            if(len1 < 0) {
                /* if error, skip frame */
                audio_pkt_size = 0;
                break;
            }

            audio_pkt_size -= len1;
            int data_size = 0;

            if(got_frame) {


                if(pVideoInfo->audioOutSpec.bufferSize > 0) {
                    data_size = pVideoInfo->audioOutSpec.bufferSize;
                } else {
                    data_size = av_samples_get_buffer_size(NULL,
                                                           pVideoInfo->pAudioCodecCtx->channels,
                                                           frame.nb_samples,
                                                           pVideoInfo->pAudioCodecCtx->sample_fmt,
                                                           1);
                }

                swr_convert(pVideoInfo->pSwrCtx, &audio_buf, MAX_AUDIO_FRAME_SIZE, (const uint8_t **)frame.data , frame.nb_samples);

#if WRITE_PCM
                fwrite(audio_buf, 1, data_size, pcmfile);
#endif
            } else {
                continue;
            }

            /* We have data, return it and come back for more later */
            return data_size;
        } // while

        if(pkt.data)
            av_free_packet(&pkt);

        pVideoInfo->audioQueue.Pop(&pkt, true);

        audio_pkt_size = pkt.size;
    }

	return -1;
}


void SDLCALL audio_callback(void *userdata, Uint8 *stream, int len) {
    VideoInfo *pVideoInfo = (VideoInfo *)userdata;
    int len1, audio_size;

    static uint8_t audio_buf[(MAX_AUDIO_FRAME_SIZE * 3) / 2];
    static unsigned int audio_buf_size = 0;
    static unsigned int audio_buf_index = 0;

    while(len > 0) {
        if(audio_buf_index >= audio_buf_size) {
            /* We have already sent all our data; get more */
            audio_size = DecodeAudioFrame(pVideoInfo, audio_buf, sizeof(audio_buf));

            if(audio_size < 0) {
                /* If error, output silence */
                audio_buf_size = 1024; // arbitrary?
                memset(audio_buf, 0, audio_buf_size);
            } else {
                audio_buf_size = audio_size;
            }

            audio_buf_index = 0;
        }

        len1 = audio_buf_size - audio_buf_index;

        if(len1 > len)
            len1 = len;

        memcpy(stream, (uint8_t *)audio_buf + audio_buf_index, len1);
        len -= len1;
        stream += len1;
        audio_buf_index += len1;
    }

}

void VideoDisplay(VideoInfo *pVideoInfo) {
    SDL_Rect rect;
    VideoPicture vp;
    float aspectRatio;

    if(!pVideoInfo->picQueue.Pop(&vp, 0))
        return;

    if(!vp.pBmp)
        return;

    if(pVideoInfo->pVideoCodecCtx->sample_aspect_ratio.num == 0) {
        aspectRatio = 0;
    } else {
        aspectRatio = av_q2d(pVideoInfo->pVideoCodecCtx->sample_aspect_ratio) * pVideoInfo->pVideoCodecCtx->width / pVideoInfo->pVideoCodecCtx->height;
    }

    if(aspectRatio <= 0.0) {
        aspectRatio = (float)pVideoInfo->pVideoCodecCtx->width / (float)pVideoInfo->pVideoCodecCtx->height;
    }

    int h = pVideoInfo->pScreen->h;
    int w = ROUND(h * aspectRatio) & -3;

    if(w > pVideoInfo->pScreen->w) {
        w = pVideoInfo->pScreen->w;
        h = ROUND(w / aspectRatio) & -3;
    }

    int x = (pVideoInfo->pScreen->w - w) / 2;
    int y = (pVideoInfo->pScreen->h - h) / 2;

    rect.x = x;
    rect.y = y;
    rect.w = w;
    rect.h = h;

	SDL_LockMutex(pVideoInfo->pScreenMutex);
	SDL_DisplayYUVOverlay(vp.pBmp, &rect);
	SDL_UnlockMutex(pVideoInfo->pScreenMutex);

	SDL_FreeYUVOverlay(vp.pBmp);
}

void VideoDisplay2(VideoInfo* pVideoInfo, VideoPicture *vp) {
	SDL_Rect rect;
	float aspectRatio;


	if(!vp->pBmp)
		return;

	if(pVideoInfo->pVideoCodecCtx->sample_aspect_ratio.num == 0) {
		aspectRatio = 0;
	} else {
		aspectRatio = av_q2d(pVideoInfo->pVideoCodecCtx->sample_aspect_ratio) * pVideoInfo->pVideoCodecCtx->width / pVideoInfo->pVideoCodecCtx->height;
	}

	if(aspectRatio <= 0.0) {
		aspectRatio = (float)pVideoInfo->pVideoCodecCtx->width / (float)pVideoInfo->pVideoCodecCtx->height;
	}

	int h = pVideoInfo->pScreen->h;
	int w = ROUND(h * aspectRatio) & -3;

	if(w > pVideoInfo->pScreen->w) {
		w = pVideoInfo->pScreen->w;
		h = ROUND(w / aspectRatio) & -3;
	}

	int x = (pVideoInfo->pScreen->w - w) / 2;
	int y = (pVideoInfo->pScreen->h - h) / 2;

	rect.x = x;
	rect.y = y;
	rect.w = w;
	rect.h = h;

	SDL_DisplayYUVOverlay(vp->pBmp, &rect);

	SDL_FreeYUVOverlay(vp->pBmp);
}

Uint32 SDLCALL SDLRefreshTimerCB(Uint32 interval, void *param) {
	SDL_Event event;
	event.type = FF_REFRESH_EVENT;
	event.user.data1 = param;

	SDL_PushEvent(&event);

	return 0; /* 0 means stop timer */
}

void ScheduleRefresh(VideoInfo *pVideoInfo, int delay) {
	SDL_AddTimer(delay, SDLRefreshTimerCB, pVideoInfo);
}



void VideoRefreshTimer(void *pUserData) {
	VideoInfo *pVideoInfo = (VideoInfo*)pUserData;
	if(pVideoInfo->pVideoStream) {
		ScheduleRefresh(pVideoInfo, 40);
		VideoDisplay(pVideoInfo);
	} else {
		ScheduleRefresh(pVideoInfo, 100);
	}
}




int VideoDecodeThread(void *arg) {
    VideoInfo *pVideoInfo = (VideoInfo *)arg;
    AVPacket pkt;
    AVFrame *pFrame = av_frame_alloc();
    int frameFinished = 0;

    pVideoInfo->pScreen = SDL_SetVideoMode(pVideoInfo->pVideoCodecCtx->width, pVideoInfo->pVideoCodecCtx->height, 0, 0);

    while(!pVideoInfo->quit) {
        if(!pVideoInfo->videoQueue.Pop(&pkt, 0))
            continue;

        avcodec_decode_video2(pVideoInfo->pVideoCodecCtx, pFrame, &frameFinished, &pkt);

        if(frameFinished > 0) {
            VideoPicture vp;
			SDL_LockMutex(pVideoInfo->pScreenMutex);
            vp.pBmp = SDL_CreateYUVOverlay(pVideoInfo->pVideoCodecCtx->width,
                                           pVideoInfo->pVideoCodecCtx->height,
                                           SDL_YV12_OVERLAY,
                                           pVideoInfo->pScreen);
			SDL_UnlockMutex(pVideoInfo->pScreenMutex);
            vp.width = pVideoInfo->pVideoCodecCtx->width;
            vp.height = pVideoInfo->pVideoCodecCtx->height;
            vp.allocated = 1;

            if(vp.pBmp) {
                SDL_LockYUVOverlay(vp.pBmp);

                AVPicture pict;
                pict.data[0] = vp.pBmp->pixels[0];
                pict.data[1] = vp.pBmp->pixels[2];
                pict.data[2] = vp.pBmp->pixels[1];

                pict.linesize[0] = vp.pBmp->pitches[0];
                pict.linesize[1] = vp.pBmp->pitches[2];
                pict.linesize[2] = vp.pBmp->pitches[1];

                sws_scale(pVideoInfo->pSwsCtx, (uint8_t const * const *)pFrame->data,
                          pFrame->linesize, 0, pVideoInfo->pVideoCodecCtx->height,
                          pict.data, pict.linesize);

                SDL_UnlockYUVOverlay(vp.pBmp);

                pVideoInfo->picQueue.Push(&vp, 1);

				//VideoDisplay(pVideoInfo);
				//VideoDisplay2(pVideoInfo, &vp);
            }
        }

        av_free_packet(&pkt);
    }

    av_frame_free(&pFrame);

    return 0;
}

bool OpenAudioStream(VideoInfo *pVideoInfo) {
    if(pVideoInfo->audioStream < 0 || pVideoInfo->audioStream >= pVideoInfo->pFmtCtx->nb_streams) {
        return false;
    }

    AVCodec *pCodec = avcodec_find_decoder(pVideoInfo->pAudioStream->codec->codec_id);

    if(!pCodec) {
        return false;
    }

    pVideoInfo->pAudioCodecCtx = avcodec_alloc_context3(pCodec);

    if(avcodec_copy_context(pVideoInfo->pAudioCodecCtx, pVideoInfo->pAudioStream->codec) != 0) {
        return false;
    }

    if(avcodec_open2(pVideoInfo->pAudioCodecCtx, pCodec, NULL) != 0) {
        return false;
    }

	pVideoInfo->audioOutSpec.sampleRate = pVideoInfo->pAudioCodecCtx->sample_rate;
    pVideoInfo->audioOutSpec.channels = av_get_channel_layout_nb_channels(pVideoInfo->audioOutSpec.channelLayout);
    pVideoInfo->audioOutSpec.samplesPerChannel = pVideoInfo->pAudioCodecCtx->frame_size;

    pVideoInfo->audioOutSpec.bufferSize = 0;

    if(pVideoInfo->audioOutSpec.samplesPerChannel > 0) {
        pVideoInfo->audioOutSpec.bufferSize = av_samples_get_buffer_size(NULL, pVideoInfo->audioOutSpec.channels,
                                              pVideoInfo->audioOutSpec.samplesPerChannel,
                                              pVideoInfo->audioOutSpec.sampleFmt,
                                              1);
    }


    int inChannleLayout = av_get_default_channel_layout(pVideoInfo->pAudioCodecCtx->channels);

    pVideoInfo->pSwrCtx = swr_alloc();
    pVideoInfo->pSwrCtx = swr_alloc_set_opts(pVideoInfo->pSwrCtx,
                          pVideoInfo->audioOutSpec.channelLayout,
                          pVideoInfo->audioOutSpec.sampleFmt,
                          pVideoInfo->audioOutSpec.sampleRate,
                          inChannleLayout,
                          pVideoInfo->pAudioCodecCtx->sample_fmt,
                          pVideoInfo->pAudioCodecCtx->sample_rate,
                          0, NULL);
    swr_init(pVideoInfo->pSwrCtx);

    return true;
}


bool OpenVideoStream(VideoInfo *pVideoInfo) {
    if(pVideoInfo->videoStream < 0 || pVideoInfo->videoStream >= pVideoInfo->pFmtCtx->nb_streams) {
        return false;
    }

    AVCodec *pCodec = avcodec_find_decoder(pVideoInfo->pVideoStream->codec->codec_id);

    if(!pCodec) {
        return false;
    }

    pVideoInfo->pVideoCodecCtx = avcodec_alloc_context3(pCodec);

    if(avcodec_copy_context(pVideoInfo->pVideoCodecCtx, pVideoInfo->pVideoStream->codec) != 0) {
        return false;
    }

    if(avcodec_open2(pVideoInfo->pVideoCodecCtx, pCodec, NULL) != 0) {
        return false;
    }

    pVideoInfo->pSwsCtx = sws_getContext(pVideoInfo->pVideoCodecCtx->width,
                                         pVideoInfo->pVideoCodecCtx->height,
                                         pVideoInfo->pVideoCodecCtx->pix_fmt,
                                         pVideoInfo->pVideoCodecCtx->width,
                                         pVideoInfo->pVideoCodecCtx->height,
                                         AV_PIX_FMT_YUV420P,
                                         SWS_BILINEAR,
                                         NULL, NULL, NULL);

    pVideoInfo->pVideoDecodeThread = SDL_CreateThread(VideoDecodeThread, pVideoInfo);

    return true;
}


int SplitPacketThread(void *arg) {
    VideoInfo *pVideoInfo = (VideoInfo *)arg;

    pVideoInfo->audioStream = -1;
    pVideoInfo->videoStream = -1;

    if(avformat_open_input(&pVideoInfo->pFmtCtx, pVideoInfo->filename, NULL, NULL) != 0) {
        return -1;
    }

    if(avformat_find_stream_info(pVideoInfo->pFmtCtx, NULL) < 0) {
        return -1;
    }

    for(unsigned int  i = 0; i < pVideoInfo->pFmtCtx->nb_streams; i++) {
        av_dump_format(pVideoInfo->pFmtCtx, i, pVideoInfo->filename, 0);
    }


    for(unsigned int  i = 0; i < pVideoInfo->pFmtCtx->nb_streams; i++) {
        if(pVideoInfo->pFmtCtx->streams[i]->codec->codec_type == AVMEDIA_TYPE_AUDIO && pVideoInfo->audioStream == -1) {
            pVideoInfo->audioStream = i;
            pVideoInfo->pAudioStream = pVideoInfo->pFmtCtx->streams[i];
        }

        if(pVideoInfo->pFmtCtx->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO && pVideoInfo->videoStream == -1) {
            pVideoInfo->videoStream = i;
            pVideoInfo->pVideoStream = pVideoInfo->pFmtCtx->streams[i];
        }
    }

    pVideoInfo->audioOutSpec.channelLayout = AV_CH_LAYOUT_STEREO;
    pVideoInfo->audioOutSpec.sampleFmt = AV_SAMPLE_FMT_S16;
    //pVideoInfo->audioOutSpec.sampleRate = 48000;
    pVideoInfo->audioOutSpec.pfnCallback = audio_callback;

    OpenAudioStream(pVideoInfo);
	OpenVideoStream(pVideoInfo);

    // SDL
    SDL_AudioSpec realSpec;
    SDL_AudioSpec wantedSpec;
    wantedSpec.callback = pVideoInfo->audioOutSpec.pfnCallback;
    wantedSpec.channels = pVideoInfo->audioOutSpec.channels;
    wantedSpec.format = AUDIO_S16SYS;
    wantedSpec.freq = pVideoInfo->audioOutSpec.sampleRate;
    wantedSpec.samples = pVideoInfo->audioOutSpec.samplesPerChannel;
    wantedSpec.silence = 0;
    wantedSpec.userdata = pVideoInfo;

    if(SDL_OpenAudio(&wantedSpec, &realSpec) < 0) {
        return -1;
    }

    SDL_PauseAudio(0);


    AVPacket pkt;

    while(true) {
        if(pVideoInfo->quit) {
            break;
        }

        if(av_read_frame(pVideoInfo->pFmtCtx, &pkt) < 0) {
            if(pVideoInfo->pFmtCtx->pb->error == 0) {
                SDL_Delay(100);
                continue;
            } else {
                break;
            }
        }

        if(pkt.stream_index == pVideoInfo->videoStream) {
            pVideoInfo->videoQueue.Push(&pkt, pkt.size);
			printf("%d\n", pVideoInfo->videoQueue.GetElementCount());
        } else if(pkt.stream_index == pVideoInfo->audioStream) {
            pVideoInfo->audioQueue.Push(&pkt, pkt.size);
        } else {
            av_free_packet(&pkt);
        }
    }

    while(!pVideoInfo->quit) {
        SDL_Delay(100);
    }

    SDL_Event event;
    event.type = FF_QUIT_EVENT;
    event.user.data1 = pVideoInfo;
    SDL_PushEvent(&event);

    return 0;
}



int main(int argc, char *argv[]) {

    VideoInfo videoInfo;
	videoInfo.audioQueue.SetMaxSize(1024);
	videoInfo.videoQueue.SetMaxSize(1024);
	videoInfo.picQueue.SetMaxSize(1024);

	videoInfo.pScreenMutex = SDL_CreateMutex();

    strcpy_s(videoInfo.filename, sizeof(videoInfo.filename), argv[1]);


    if(SDL_Init(SDL_INIT_AUDIO | SDL_INIT_VIDEO | SDL_INIT_TIMER))
        PRINT_RETURN_1("SDL_Init");

    av_register_all();

    ScheduleRefresh(&videoInfo, 40);

    videoInfo.pSplitPacketThread = SDL_CreateThread(SplitPacketThread, &videoInfo);

    SDL_Event   event;

    while(true) {
        SDL_WaitEvent(&event);

        switch(event.type) {
            case FF_QUIT_EVENT:
            case SDL_QUIT:
                videoInfo.quit = 1;
                SDL_Quit();
                return 0;
                break;

            case FF_REFRESH_EVENT:
                VideoRefreshTimer(event.user.data1);
                break;

            default:
                break;
        }
    }


    return 0;
}